# -*- coding:utf-8 -*-
"""
针对sinahotnews3数据集！！

用来统计每个股票每天有多少新闻，新闻数目就可以看作是今天该新闻的热度

对于每一条新闻，获取该新闻对应的股票列表，发布日期
1.对于日期建立dict，可以快速定位到date
2.在date下面对stock建立dict，可以快速定位到stock
3.结合上述两步，每次对相应日期相应股票进行+1操作，最终按照日期来记录每个stock的出现次数即可
4.只需要遍历一次数据库，便捷高效
author: yinzm
"""
import pymongo
from datetime import datetime
import pandas as pd
import os
import pickle

mongoClient = pymongo.MongoClient("localhost", 27017)
db = mongoClient.sinafinance
collection = db.sinahotnews3 # 注意数据库别写错了

def dateList(beginDate, endDate):
    # 产生所有的日期, beginDate: str, endDate: str
    dateSimple = [datetime.strftime(x,"%Y%m%d") for x in list(pd.date_range(start=beginDate, end=endDate))]
    return dateSimple

allDateList = dateList("20130101", "20180331")

# 将所有股票加载进来
pklPath = '../DataPreprocess/ExtractWindInfo/data'
stockCodeList = []
with open(os.path.join(pklPath, 'stockList.csv'), 'r', encoding='utf-8') as f:
    for stock in f.readlines():
        stockCode = stock.strip()[:-3]
        stockCodeList.append(stockCode)

# 初始化
dateDict = dict()
for d in allDateList:
    dateDict[d] = dict()
    for s in stockCodeList:
        dateDict[d][s] = 0

# 开始遍历数据库
cnt = 0
for item in collection.find():
    date = str(item['date'])
    stockList = item['stock']
    for stock in stockList.strip().split(','):
        if stock.isdigit() and stock in stockCodeList and date in allDateList:
            dateDict[date][stock] += 1
    cnt += 1
    if cnt%100 == 0:
        print(cnt)

# 保存结果
dateFile = 'E:\\yinzm\\nlp\\data\\Wu_HotFactorData'
for d in allDateList:
    outputFile = os.path.join(dateFile, d+'.csv')
    with open(outputFile, 'w', encoding='utf-8') as f:
        for stock in stockCodeList:
            f.write("%s,%s\n" % (stock,dateDict[d][stock]))

print("==============over===================")
